# Wei Jiang

from twisted.internet.protocol import Factory
from twisted.internet import reactor, protocol
from twisted.protocols.basic import LineReceiver
from twisted.internet.endpoints import TCP4ClientEndpoint
from ctypes import *
import pickle
import random

# single result 
class RESULT(Structure):
    _fields_ = [("score",c_double),
                ("docid",c_ulong),
                ("url",c_char_p)]

# pointer to results list (in C++)
Results_Table = POINTER(RESULT)



# proxy query processor object
class Foo(object):

    def __init__(self, protocol):
        # the lib we are going to load.
        self.lib = CDLL('/data1/team/weijiang/machine-learning-project-related/programs/trunk/polyIRIndexer/irtk.so')
        
        #self.foo = self.lib.Foo_new(top_k)
        self.foo = self.lib.Foo_new()
        
        self.protocol = protocol
        
        #print "self.protocol:",self.protocol

    def __del__(self):
        self.lib.Foo_delete(self.foo)

    # push a query for processing
    def push(self, queryId, channel, query, top_k):
        self.lib.Foo_push(self.foo, queryId, channel, query, top_k, CMPFUNC(self.callback))
        
    def report(self):
        self.lib.Foo_report(self.foo)

    # self.callback prototype for results
    #   returns int 0 = done with results (mem free() on the CPP side)
    def callback(self, qid, channel, numResults, results_ptr):
        print "in callback function"
        if qid != -1:
            print "Got ", numResults, "results for (queryID,channelID):", "(",qid,",",channel,")"
            Results = cast(results_ptr,POINTER(RESULT))
            topReturnDocIDForQuery = []
            #For each query, will have a new tuple.
            for item in Results:
                # list is null terminated
                if item.url is None:
                    break
                # do something with the results
                
                #current version
                outputLine = "Score: " + "{0:.6f}".format( float(item.score)) + "    "+ "DocID: " + str(item.docid) +"    "+ "URL: " + item.url
                print outputLine
                
                # currently, there are no url involved. 
                # I will have to fix that in a later case.
                
                topReturnDocIDForQuery.append( (int(item.docid), float(item.score) ,str(item.url) ))
                #topReturnDocIDForQuery.append( int(item.docid) )
            
            newTuple = (qid,topReturnDocIDForQuery)   
            self.protocol.queryResultDict[channel].append(newTuple)
        else:
            #The ending mark of the query.
            # now change parameters to be dynamic.
            # (1)check the number of channels are there.
            # (2)check the number of results are there.
            if self.whetherAnswersAreReadyToSend():
                print self.protocol.queryResultDict
                print "length to send for the outgoing queryDict:",len(pickle.dumps(self.protocol.queryResultDict))
                self.protocol.sendLine( pickle.dumps(self.protocol.queryResultDict) )
            
            pass
        return 0
    
    def whetherAnswersAreReadyToSend(self):
        flag = True
        
        '''
        print "test1 in"
        print "len( self.protocol.queryResultDict ):",len( self.protocol.queryResultDict )
        print "len(self.protocol.queryDict):",len(self.protocol.queryDict)
        print "test1 out"
        '''
        
        if len( self.protocol.queryResultDict ) == len(self.protocol.queryDict):
            pass
        else:
            flag = False
            return flag
        
        for channelNumber in self.protocol.queryDict:
            if len(self.protocol.queryResultDict[channelNumber]) == len(self.protocol.queryDict[channelNumber]):
                '''
                print "test2 in"
                print "channelNumber:",channelNumber
                print "len(self.protocol.queryResultDict[channelNumber])",len(self.protocol.queryResultDict[channelNumber])
                print "len(self.protocol.queryDict[channelNumber])",len(self.protocol.queryDict[channelNumber])
                print "test2 out"
                '''
                
                pass
            else:
                flag = False
                return flag
        
        return flag

class slaveClientCommand(LineReceiver):
    def __init__(self,factory):
        self.slaveClientStateCode = None
        self.slaveID = factory.slaveID
        
        #self.clientCommandInitSequence = [(2,1),(1,1,3),(4,1),(2,1),(0,1,2),(0,1,3),(2,1),(3,1)]
        
        if self.slaveID == 0:
            self.clientCommandInitSequence = [(8,1,"apple"),(7,1),(3,1)]
        elif self.slaveID == 1:
            self.clientCommandInitSequence = [(7,1),(2,1),(0,1,2),(0,1,3),(3,1)]
        elif self.slaveID == 2:
            self.clientCommandInitSequence = [(7,1),(2,1),(0,1,3),(3,1)]
        elif self.slaveID == 9999:
            #for testing: single query in an interactive mode.
            self.clientCommandInitSequence = [(7,1),(3,1)]
        
        # TODO: The server side won't decrease the # of connections.
        #self.clientCommandInitSequence = [(5,1)]
        self.brokerServerCommands = factory.brokerServerCommands
        self.slaveClientStates = factory.slaveClientStates
        self.typeCode = factory.typeCode
        self.factory = factory
        self.currentSentBrokerCommandCode = -1
        self.registedChannels = []
        self.slaveClientStateCode = 0
        
        if self.slaveID != 0 and self.slaveID != 9999:
            self.loadLocalSearchEngine()
    
  
    def loadLocalSearchEngine(self):
        print "Init local search engine black box..."
        #input1000QueryFileHandler = open('/data1/team/weijiang/machine-learning-project-related/programs/polyIRToolkit/polyIRIndexer/10query','r')

        # Init the CPP query processor proxy
        self.localSearchEngineBlackBox = Foo(self)
        
        
    def sendNextCommandToServer(self):
        print "    slave client state:",self.slaveClientStateCode,self.slaveClientStates[self.slaveClientStateCode]
        if len(self.clientCommandInitSequence) != 0:
            tuple = self.clientCommandInitSequence.pop()
            if len(tuple) == 2:
                (self.currentSentBrokerCommandCode,mode) = tuple
                if self.currentSentBrokerCommandCode in self.brokerServerCommands:
                    outgoingLine = str(self.currentSentBrokerCommandCode) 
                    outgoingLineVerbose = str(self.currentSentBrokerCommandCode) + " " + self.brokerServerCommands[self.currentSentBrokerCommandCode]
                    print "        ---->Sent:",outgoingLineVerbose
                    self.sendLine(outgoingLine)
                    if mode == 0:
                        pass
                    elif mode == 1:
                        self.setRawMode()
                else:
                    print "Not Supported Command Code:",self.currentSentBrokerCommandCode
                    self.slaveClientStateCode = 1
                    #I am not sure the following code is right.
                    self.transport.loseConnection()
                    
            elif len(tuple) == 3:
                (self.currentSentBrokerCommandCode,mode,content) = tuple
                if self.currentSentBrokerCommandCode in self.brokerServerCommands:
                    outgoingLine = str(self.currentSentBrokerCommandCode) + " " + str(content)
                    outgoingLineVerbose = str(self.currentSentBrokerCommandCode) + " "+ self.brokerServerCommands[self.currentSentBrokerCommandCode] + " "+ str(content)
                    print "        ---->Sent:",outgoingLineVerbose
                    
                    # the code means SubscribeChannels, so add the channel number from the list.
                    if self.currentSentBrokerCommandCode == 0:
                        assert content not in self.registedChannels
                        self.registedChannels.append(content)
                    
                    # the code means UnSubscribeChannels, so delete the channel number from the list.
                    if self.currentSentBrokerCommandCode == 1:
                        assert content in self.registedChannels
                        self.registedChannels.remove(content)                    
                    
                    self.sendLine(outgoingLine)
                    if mode == 0:
                        pass
                    elif mode == 1:
                        self.setRawMode()
                else:
                    print "Not Supported Command Code:",self.currentSentBrokerCommandCode
                    self.slaveClientStateCode = 1
                    #I am not sure the following code is right.
                    self.transport.loseConnection()
                    
        else:
            if self.slaveID != 9999:
                exitReason = "Client Init...Done. Waiting for commands from server..."
                print "        ---->",exitReason
            else:
                self.slaveClientStateCode = 4
                queryContentEntered = raw_input('Distributed Search:')
                queryContentEntered = queryContentEntered.lower()
                self.clientCommandInitSequence.append( (8,1,queryContentEntered) )
                self.sendNextCommandToServer()

    
    def printWatingForCommandsFromServerMessageLine(self):
        exitReason = "Client Operation...Done. Waiting for commands from server..."
        print "        ---->",exitReason
    
    def connectionMade(self):
        self.factory.numCommandProtocols = self.factory.numCommandProtocols + 1
        #print "There are currently",self.factory.numCommandProtocols,"open command connection"
        # init the command
        
        self.sendNextCommandToServer()
        
        

    def lineReceived(self, line):
        print "lineReceived:",line
        print "there is nothing currently,here."
    
    def rawDataReceived(self,data):
        #print "rawDataReceived:",data
        #print "self.slaveClientStateCode:",self.slaveClientStateCode
        if data.strip() == "FAIL":
            print "FAIL."
            # This FAIL will lead the client to close its operation right now.
            self.slaveClientStateCode = 1
            self.transport.loseConnection()
        else:
            if self.slaveClientStateCode != 3:
                if self.currentSentBrokerCommandCode == 0:
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 1:
                    pass
                elif self.currentSentBrokerCommandCode == 2:
                    endpointsWithChannelsDict = pickle.loads(data)
                    for endpoint in endpointsWithChannelsDict:
                        print "            ",endpoint,endpointsWithChannelsDict[endpoint]
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 3:
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 4:
                    print pickle.loads(data)
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 5:
                    ingoingQueryDict = pickle.loads(data)
                    self.handle_IngoingQueriesFromServer(ingoingQueryDict)
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 6:
                    self.slaveClientStateCode = 2
                elif self.currentSentBrokerCommandCode == 7:
                    self.slaveClientStateCode = 3
                elif self.currentSentBrokerCommandCode == 8:
                    self.slaveClientStateCode = 4                
                else:
                    print "Unknown command code:",self.currentSentBrokerCommandCode
                self.sendNextCommandToServer()
            else:
                ingoingQueryDict = pickle.loads(data)
                self.handle_IngoingQueriesFromServer(ingoingQueryDict)
                
    
    def handle_IngoingQueriesFromServer(self, queryDict):
        print "handle_IngoingQueriesFromServer"
        self.queryResultDict = {}
        self.queryDict = queryDict
        for channelNumber in self.queryDict:
            print "channelNumber:",channelNumber
            self.queryResultDict[channelNumber] = []
            
            # Push all the queies into the local search engine balck box with the ending mark as: queryNumber:-1 ,channelNumber:the same,queryContent:""
            for tuple in self.queryDict[channelNumber]:
                #print tuple
                (queryNumber,queryContent, top_k) = tuple
                
                self.localSearchEngineBlackBox.push(queryNumber,channelNumber,queryContent, top_k)
        
            #Push the ending mark as well.
            self.localSearchEngineBlackBox.push(-1,channelNumber,"", 0)
        
        
        
        self.localSearchEngineBlackBox.report()
        self.printWatingForCommandsFromServerMessageLine()

    def connectionLost(self, reason):
        self.factory.numCommandProtocols = self.factory.numCommandProtocols - 1
        #print "There are currently",self.factory.numCommandProtocols,"open command connection"
        print "    slave client state:",self.slaveClientStateCode,self.slaveClientStates[self.slaveClientStateCode]
        reactor.stop()

    def testFunction(self):
        print "This means the good sign."

class slaveClientFactory(Factory):
    def __init__(self, typeCode,slaveID):
        self.typeCode = typeCode
        self.slaveID = slaveID
        self.numCommandProtocols = 0
        self.numDataProtocols = 0
        self.brokerServerCommands = {0:"SubscribeChannels", 1:"UnsubscribeChannels", 2:"GetStatus",3:"RegisterMyself",4:"GetChannels",5:"QueryMe",6:"QueryOthers",7:"IAMReady",8:"QueryOthersOneQueryPerTime"}
        self.slaveClientStates = {0:"INIT",1:"END",2:"REGISTERED",3:"READY",4:"READYTOSEARCHOTHERS"}
        

    def buildProtocol(self, addr):
        
        if self.typeCode == 0:
            return slaveClientCommand(self)
        elif self.typeCode == 1:
            return slaveClientData(self)



def gotProtocol(p):
    if p.typeCode == 0:
        commandProtocolInstanceReference = p
    elif p.typeCode == 1:
        print "data:",p
        dataProtocolInstanceReference = p        
    else:
        print "typeCode NOT recognized."

#this slave factory can generate 2 kinds of products based on different type codes.
#typeCode: 0 for control port, 1 for data port
print "Program begins."
commandProtocolInstanceReference = None
dataProtocolInstanceReference = None

slaveID = -1
while slaveID != 0 and slaveID != 1 and slaveID != 2 and slaveID != 9999:
    try:
        slaveID = int( raw_input('slaveID(0, 1, 2 or 9999):') )
    except:
        pass


#default setting for server and the port.
controlPort = 10877
serverName = "localhost"

controlPoint = TCP4ClientEndpoint(reactor,serverName,controlPort)
commandPortFactory = slaveClientFactory(0,slaveID)
deferredForCommandPort = controlPoint.connect(commandPortFactory)
deferredForCommandPort.addCallback(gotProtocol)

CMPFUNC = CFUNCTYPE(c_int,c_int, c_int, c_int, POINTER(RESULT))


reactor.run()
print "Program end."
